require './n_gram'

class NGramCounter
  attr_accessor :grams, :voc
  attr_reader :arn, :total

  def initialize n, voc = []
    @arn = n
    @grams = []
    @voc = voc
    @total = 0
  end

  def voc_index word
    return @voc.index(word)+1 if @voc.include? word
    @voc.push word
    @voc.size
  end

  def << ngram
    ngram = ngram.to_a if ngram.is_a? NGram
    root = @grams
    ngram[0..-2].each { |gram|
      voc_id = voc_index gram
      root[voc_id] = [0] unless root[voc_id]
      root = root[voc_id]
    }
    voc_id = voc_index ngram[-1]
    if ngram.size == @arn
      root[voc_id] = 0 unless root[voc_id]
      root[voc_id] += 1
    else
      root[voc_id] = [] unless root[voc_id]
      root[voc_id][0] = 0 unless root[voc_id][0]
      root[voc_id][0] += 1
    end
    @total += 1
  end

  def core_print root = nil, level = 1, not_less_than = 0, prefix = ""
    (1..root.size-1).each { |i|
      next unless root[i]
      nx_prefix = prefix + @voc[i-1] + " "
      if level == 1
        count = root[i].is_a?(Array) ? root[i][0] : root[i]
        next unless count && count >= not_less_than
        puts nx_prefix + count.to_s
        next
      end
      core_print root[i], level-1, not_less_than, nx_prefix
    }
  end

  def count words
    words = words.to_a if words.is_a? NGram
    root = @grams
    words.each { |word|
      return 0 unless root
      root = root[voc_index word]
    }
    return 0 unless root
    return root[0] if root.is_a? Array
    return root if root.is_a? Integer
  end

  def print not_less_than = 0
    (1..@arn).each { |ns|
      puts "== "+ns.to_s
      core_print @grams, ns, not_less_than
    }
  end

  def self.test
    cc = NGramCounter.new 3
    cc << %w[dima like beer]
    cc << %w[dima]
    cc << %w[dima hates acid]
    cc << %w[dima like beer]
    cc << %w[acid]
    cc << %w[dima like]
    cc << %w[acid]
    t1 = cc.count %w[dima like beer]
    t2 = cc.count %w[dima like]
    t3 = cc.count %w[acid]
    t4 = cc.count %w[quake]
    t5 = cc.count %w[dima hates]
    raise "counter not work" unless [t1,t2,t3,t4,t5] == [2,1,2,0,0]
  end
end
